"""
-*- coding: utf-8 -*-
@File  : generate_videos.py
@author: ZhenyuYang
@Time  : 2023/03/31 10:38
"""
import os
import cv2

raw_frame_path = "../data/flows_raw/"
dataset = ["casme2", "smic", "samm"]
labels = ["negative", "positive", "surprise"]
if not os.path.exists("../data/flow_videos"):
    os.mkdir("../data/flow_videos")
for label in labels:
    if not os.path.exists("../data/flow_videos/{}".format(label)):
        os.mkdir("../data/flow_videos/{}".format(label))
    raw_frame_class_path = "{}/{}".format(raw_frame_path, label)
    folder_list = os.listdir(raw_frame_class_path)
    for folder in folder_list:
        folder_dataset = folder.split("-")[0]
        if folder_dataset not in dataset:
            continue
        src_path = "{}/{}".format(raw_frame_class_path, folder)
        dst_path = "../data/flow_videos/{}/{}.avi".format(label, folder)
        frame_list = os.listdir(src_path)
        frame = cv2.imread("{}/{}".format(src_path, frame_list[0]))
        video_writer = cv2.VideoWriter(dst_path, cv2.VideoWriter_fourcc(*"MJPG"), 30, (frame.shape[1], frame.shape[0]))
        for frame_index in frame_list:
            frame = cv2.imread("{}/{}".format(src_path, frame_index))
            frame = cv2.resize(frame, (frame.shape[1], frame.shape[0]))
            video_writer.write(frame)
        video_writer.release()
        print(dst_path)
